/** @file test_admm.h
 *
 *  @author Dongryeol Lee (dongryel@cc.gatech.edu)
 */

#ifndef MLPACK_ADMM_TEST_ADMM_H
#define MLPACK_ADMM_TEST_ADMM_H

#include <boost/test/unit_test.hpp>
#include "core/gnp/dualtree_dfs_dev.h"
#include "core/parallel/random_dataset_generator.h"
#include "core/tree/gen_metric_tree.h"
#include "core/math/math_lib.h"
#include "mlpack/admm/admm_dev.h"
#include "mlpack/series_expansion/kernel_aux.h"

namespace mlpack {
namespace admm {
namespace test_admm {
extern int num_dimensions_;
extern int num_points_;
extern tbb::task_scheduler_init admm_scheduler_;
}

class TestAdmm {
  private:

  public:

    int StressTest(
      boost::mpi::communicator &world, const std::string &loss_function,
      const std::string &regularization) {

      std::string references_in("random_references.csv");
      std::string labels_in("random_labels.csv");
      std::string original_references_in = references_in;
      std::string original_labels_in = labels_in;

      // The list of arguments.
      std::vector< std::string > args;

      // Push in the reference dataset name and the random label set.
      std::stringstream reference_file_name_sstr;
      reference_file_name_sstr << references_in << world.rank();
      references_in = reference_file_name_sstr.str();
      std::stringstream label_file_name_sstr;
      label_file_name_sstr << labels_in << world.rank();
      labels_in = label_file_name_sstr.str();
      args.push_back(std::string("--references_in=") + references_in);
      args.push_back(std::string("--reference_labels_in=") + labels_in);

      // Push in the loss function type and the regularization type.
      args.push_back(std::string("--loss_in=") + loss_function);
      args.push_back(std::string("--regularization_in=") + regularization);

      // Generate the random dataset and save it.
      mlpack::admm::ADMMTableType random_table;
      core::parallel::RandomDatasetGenerator::Generate(
        world,
        mlpack::admm::test_admm::num_dimensions_,
        mlpack::admm::test_admm::num_points_, std::string("none"),
        std::string("uniform"), true, &random_table);

      // Overwrite the labels with either +1 or -1.
      printf(
        "I am going with %d on %d\n", random_table.n_entries(), world.rank());
      for(int i = 0; i < random_table.n_entries(); i++) {
        arma::vec point;
        arma::vec weight;
        random_table.get(i, &point, &weight);
        int random_integer = core::math::RandInt(0, 2);
        weight[0] = (random_integer == 0) ? -1 : 1 ;
      }

      // Save the table.
      random_table.Save(references_in, &labels_in);
      world.barrier();

      // Parse the ADMM arguments.
      mlpack::admm::ADMMArguments admm_arguments;
      boost::program_options::variables_map vm;
      mlpack::admm::ADMMArgumentParser::ConstructBoostVariableMap(
        world, args, &vm);
      mlpack::admm::ADMMArgumentParser::ParseArguments(
        world, vm, &admm_arguments);

      // Run the driver.
      mlpack::admm::ADMM admm_instance;
      admm_instance.Init(world, admm_arguments);
      mlpack::admm::ADMMResult admm_result;
      admm_instance.Compute(admm_arguments, &admm_result);

      // Run the predictions on the training set as the test set and
      // export to the files.
      admm_result.Predict(* admm_arguments.reference_table_);
      admm_result.Export(
        admm_arguments.model_out_, &admm_arguments.predictions_out_);

      // Now read in the table from all of the files to construct a
      // mega-table.
      boost::mpi::communicator world_alone = world;
      int rank_is_zero = (world.rank() == 0) ? 0 : 1;
      world_alone = world_alone.split(rank_is_zero);
      if(world.rank() == 0) {
        std::vector< std::string > reference_file_names;
        std::vector< std::string > reference_label_file_names;
        for(int i = 0; i < world.size(); i++) {
          std::stringstream tmp_reference_file_name_sstr;
          tmp_reference_file_name_sstr <<  original_references_in << i;
          std::string tmp_references_in = tmp_reference_file_name_sstr.str();
          std::stringstream tmp_label_file_name_sstr;
          tmp_label_file_name_sstr << original_labels_in << i;
          std::string tmp_labels_in = label_file_name_sstr.str();
          reference_file_names.push_back(tmp_references_in);
          reference_label_file_names.push_back(tmp_labels_in);
        }
        mlpack::admm::ADMMTableType combined_table;

        combined_table.Init(
          world_alone, reference_file_names, world_alone.size(),
          & reference_label_file_names);

        // Save the table.
        std::string combined_references_in("random_combined_references.csv");
        std::string combined_labels_in("random_combined_labels.csv");
        combined_table.Save(
          combined_references_in, & combined_labels_in);

        // Parse the ADMM arguments.
        std::vector< std::string > naive_args;
        naive_args.push_back(
          std::string("--references_in=") + combined_references_in);
        naive_args.push_back(
          std::string("--reference_labels_in=") + combined_labels_in);
        naive_args.push_back(
          std::string("--loss_in=") + loss_function);
        naive_args.push_back(
          std::string("--regularization_in=") + regularization);
        mlpack::admm::ADMMArguments naive_admm_arguments;
        boost::program_options::variables_map naive_vm;
        mlpack::admm::ADMMArgumentParser::ConstructBoostVariableMap(
          world_alone, naive_args, &naive_vm);
        mlpack::admm::ADMMArgumentParser::ParseArguments(
          world_alone, naive_vm, &naive_admm_arguments);

        // Run the driver.
        mlpack::admm::ADMM naive_admm_instance;
        naive_admm_instance.Init(world_alone, naive_admm_arguments);
        mlpack::admm::ADMMResult naive_admm_result;
        naive_admm_instance.Compute(naive_admm_arguments, &naive_admm_result);
      }

      return 0;
    }

    int StressTestMain(boost::mpi::communicator &world) {

      for(int i = 0; i < 1; i++) {

        // Randomly choose the number of dimensions and the points.
        mlpack::admm::test_admm::num_dimensions_ = core::math::RandInt(2, 5);
        boost::mpi::broadcast(world,  mlpack::admm::test_admm::num_dimensions_, 0);

        mlpack::admm::test_admm::num_points_ = core::math::RandInt(2000, 5100);
        printf("Process %d generating %d points...\n", world.rank(),
               mlpack::admm::test_admm::num_points_);

        for(int j = 0; j < 1; j++) {
          switch(j) {

              // L2 SVM with L2 regularization.
            case 0:
              StressTest(
                world, std::string("squared_hinge_loss"), std::string("l2"));
              break;

              // L1 SVM with L2 regularization.
            case 1:
              StressTest(
                world, std::string("hinge_loss"), std::string("l2"));
              break;
          }
        }
      }
      return 0;
    }

};
}
}

#endif
